#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
去水印网络训练脚本
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import lpips
import numpy as np
from PIL import Image
import time
from tqdm import tqdm
import yaml
from easydict import EasyDict
import matplotlib.pyplot as plt

# 导入网络和SepMark
from watermark_removel import LightweightWatermarkRemover, UltraLightWatermarkRemover
from network.Dual_Mark import Network as SepMarkNetwork
from watermark_embding import WatermarkEmbedder


class WatermarkedDataset(Dataset):
    """水印图像数据集"""

    def __init__(self, original_dir, watermarked_dir, transform=None, max_samples=None):
        self.original_dir = original_dir
        self.watermarked_dir = watermarked_dir
        self.transform = transform

        # 获取图像文件列表
        self.image_files = []
        for file in os.listdir(watermarked_dir):
            if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                original_path = os.path.join(original_dir, file)
                watermarked_path = os.path.join(watermarked_dir, file)
                if os.path.exists(original_path):
                    self.image_files.append(file)

        if max_samples:
            self.image_files = self.image_files[:max_samples]

        print(f"数据集大小: {len(self.image_files)}")

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        file = self.image_files[idx]

        # 加载原图和水印图
        original_path = os.path.join(self.original_dir, file)
        watermarked_path = os.path.join(self.watermarked_dir, file)

        original = Image.open(original_path).convert('RGB')
        watermarked = Image.open(watermarked_path).convert('RGB')

        if self.transform:
            original = self.transform(original)
            watermarked = self.transform(watermarked)

        return watermarked, original


class WatermarkRemovalTrainer:
    def __init__(self, config_path='config/removal_config.yaml'):
        """初始化训练器"""
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {self.device}")

        # 加载配置
        self.load_config(config_path)

        # 初始化动态权重系统
        self.setup_adaptive_weights()

        # 初始化网络
        self.setup_networks()

        # 设置损失函数
        self.setup_losses()

        # 设置优化器
        self.setup_optimizers()

    def load_config(self, config_path):
        """加载配置文件"""
        default_config = {
            'model': {
                'type': 'lightweight',
                'base_channels': 32
            },
            'training': {
                'batch_size': 4,
                'learning_rate': 1e-4,
                'num_epochs': 100,
                'max_samples': 3000
            },
            'loss_weights': {
                'mse': 1.0,
                'lpips': 1.0,
                'adversarial': 0.5,
                'ssim': 0.5
            },
            'adaptive_weights': {
                'quality_threshold': 0.9,
                'quality_weight_boost': 2.0,
                'watermark_weight_boost': 1.5,
                'quality_metric': 'ssim',
                'adjustment_patience': 5,
                'min_loss_threshold': 0.01,
                'max_weight': 10.0,
                'min_weight': 0.01
            },
            'data': {
                'original_dir': 'D:/deepfake/CelebAMask-HQ/CelebAMask-HQ/CelebA-HQ-img',
                'watermarked_dir': 'D:/deepfake/CelebAMask-HQ/CelebAMask-HQ/mark/watermarked'
            },
            'sepmark': {
                'config_path': 'cfg/test_DualMark.yaml',
                'watermark_seed': 42
            },
            'monitoring': {
                'print_batch_interval': 100,
                'save_checkpoint_interval': 20,
                'plot_history_interval': 20
            },
            'data_loading': {
                'num_workers': 4,
                'pin_memory': True,
                'shuffle': True
            }
        }

        # 深度合并配置字典
        def deep_merge(default, loaded):
            """深度合并两个字典，loaded覆盖default"""
            result = default.copy()
            for key, value in loaded.items():
                if key in result and isinstance(result[key], dict) and isinstance(value, dict):
                    result[key] = deep_merge(result[key], value)
                else:
                    result[key] = value
            return result

        if os.path.exists(config_path):
            with open(config_path, 'r', encoding='utf-8') as f:
                loaded_config = yaml.load(f, Loader=yaml.SafeLoader)
                if loaded_config is None:
                    loaded_config = {}
                
                # 合并配置，确保所有必需的键都存在
                merged_config = deep_merge(default_config, loaded_config)
                config = EasyDict(merged_config)
                
                # 如果加载的配置缺少某些键，更新配置文件
                if set(loaded_config.keys()) != set(default_config.keys()):
                    print(f"更新配置文件: {config_path}")
                    with open(config_path, 'w') as f:
                        yaml.dump(merged_config, f, default_flow_style=False)
        else:
            config = EasyDict(default_config)
            # 创建默认配置文件
            os.makedirs(os.path.dirname(config_path), exist_ok=True)
            with open(config_path, 'w') as f:
                yaml.dump(default_config, f, default_flow_style=False)
            print(f"创建默认配置文件: {config_path}")

        self.config = config
        print("配置加载完成")

    def setup_networks(self):
        """设置网络"""
        # 初始化去水印网络
        if self.config.model.type == 'lightweight':
            self.removal_net = LightweightWatermarkRemover(
                base_channels=self.config.model.base_channels
            ).to(self.device)
        else:
            self.removal_net = UltraLightWatermarkRemover(
                base_channels=self.config.model.base_channels
            ).to(self.device)

        print(f"去水印网络参数量: {sum(p.numel() for p in self.removal_net.parameters()):,}")

        # 初始化SepMark网络（用于评估）
        self.setup_sepmark()

    def setup_sepmark(self):
        """设置SepMark网络"""
        try:
            # 创建SepMark嵌入器用于获取网络结构
            embedder = WatermarkEmbedder(self.config.sepmark.config_path)
            self.sepmark_net = embedder.network.encoder_decoder
            
            # 确保SepMark网络在GPU上
            self.sepmark_net = self.sepmark_net.to(self.device)
            self.sepmark_net.eval()

            # 生成原始水印消息并转移到GPU
            self.original_message = embedder.generate_watermark_message(
                self.config.sepmark.watermark_seed
            )
            # 确保消息在GPU上
            self.original_message = self.original_message.to(self.device)
            
            print("SepMark网络加载成功")

        except Exception as e:
            print(f"警告: SepMark网络加载失败: {e}")
            self.sepmark_net = None
            self.original_message = None

    def setup_losses(self):
        """设置损失函数"""
        self.mse_loss = nn.MSELoss()
        self.l1_loss = nn.L1Loss()

        # LPIPS感知损失
        self.lpips_loss = lpips.LPIPS(net='alex').to(self.device)

        print("损失函数设置完成")

    def setup_optimizers(self):
        """设置优化器"""
        self.optimizer = optim.Adam(
            self.removal_net.parameters(),
            lr=float(self.config.training.learning_rate),
            betas=(0.9, 0.999)
        )

        # 学习率调度器
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=self.config.training.num_epochs
        )
    
    def setup_adaptive_weights(self):
        """设置自适应权重系统"""
        # 动态权重配置（从配置文件读取）
        self.adaptive_config = {
            'quality_threshold': self.config.adaptive_weights.quality_threshold,
            'quality_weight_boost': self.config.adaptive_weights.quality_weight_boost,
            'watermark_weight_boost': self.config.adaptive_weights.watermark_weight_boost,
            'quality_metric': self.config.adaptive_weights.quality_metric,
            'adjustment_patience': self.config.adaptive_weights.adjustment_patience,
            'min_loss_threshold': self.config.adaptive_weights.min_loss_threshold,
            'max_weight': self.config.adaptive_weights.max_weight,
            'min_weight': self.config.adaptive_weights.min_weight,
        }
        
        # 权重历史记录
        self.weight_history = {
            'quality_weights': [],
            'watermark_weights': [],
            'quality_scores': [],
            'total_losses': []
        }
        
        # 动态权重
        self.dynamic_weights = {
            'mse': self.config.loss_weights.mse,
            'lpips': self.config.loss_weights.lpips,
            'ssim': self.config.loss_weights.ssim,
            'adversarial': self.config.loss_weights.adversarial
        }
        
        # 监控变量
        self.low_quality_count = 0
        self.stable_quality_count = 0
        self.last_quality_score = 0.0
        
        print("自适应权重系统初始化完成")
    
    def adjust_weights_based_on_quality(self, quality_score, total_loss):
        """根据图像质量动态调整权重"""
        # 记录历史
        self.weight_history['quality_scores'].append(quality_score)
        self.weight_history['total_losses'].append(total_loss)
        
        # 计算质量权重（MSE + LPIPS + SSIM）
        quality_weight = self.dynamic_weights['mse'] + self.dynamic_weights['lpips'] + self.dynamic_weights['ssim']
        watermark_weight = self.dynamic_weights['adversarial']
        
        # 权重调整逻辑
        if quality_score < self.adaptive_config['quality_threshold']:
            # 图像质量低，增加质量权重
            self.low_quality_count += 1
            self.stable_quality_count = 0
            
            if self.low_quality_count >= self.adaptive_config['adjustment_patience']:
                # 提升质量相关权重
                boost_factor = self.adaptive_config['quality_weight_boost']
                self.dynamic_weights['mse'] *= boost_factor
                self.dynamic_weights['lpips'] *= boost_factor
                self.dynamic_weights['ssim'] *= boost_factor
                
                # 降低水印权重，避免过度对抗
                self.dynamic_weights['adversarial'] *= 0.8
                
                print(f"⚠️  质量低于阈值 ({quality_score:.3f} < {self.adaptive_config['quality_threshold']})，增加质量权重")
                print(f"   新权重 - MSE: {self.dynamic_weights['mse']:.3f}, LPIPS: {self.dynamic_weights['lpips']:.3f}, SSIM: {self.dynamic_weights['ssim']:.3f}")
                
                self.low_quality_count = 0
                
        else:
            # 图像质量达标，可以考虑增加水印消除权重
            self.stable_quality_count += 1
            self.low_quality_count = 0
            
            if (self.stable_quality_count >= self.adaptive_config['adjustment_patience'] and 
                total_loss > self.adaptive_config['min_loss_threshold']):
                
                # 在质量稳定的情况下，增加水印消除权重
                watermark_boost = self.adaptive_config['watermark_weight_boost']
                self.dynamic_weights['adversarial'] *= watermark_boost
                
                print(f"✓ 质量稳定 ({quality_score:.3f} >= {self.adaptive_config['quality_threshold']})，增加水印消除权重")
                print(f"   新权重 - Adversarial: {self.dynamic_weights['adversarial']:.3f}")
                
                self.stable_quality_count = 0
        
        # 记录权重历史
        current_quality_weight = self.dynamic_weights['mse'] + self.dynamic_weights['lpips'] + self.dynamic_weights['ssim']
        self.weight_history['quality_weights'].append(current_quality_weight)
        self.weight_history['watermark_weights'].append(self.dynamic_weights['adversarial'])
        
        # 防止权重过度增长
        self.clamp_weights()
        
        self.last_quality_score = quality_score
        
        return self.dynamic_weights.copy()
    
    def clamp_weights(self):
        """限制权重范围，防止过度增长"""
        max_weight = self.adaptive_config['max_weight']
        min_weight = self.adaptive_config['min_weight']
        
        for key in self.dynamic_weights:
            self.dynamic_weights[key] = max(min_weight, min(max_weight, self.dynamic_weights[key]))
    
    def calculate_image_quality_score(self, removed, original):
        """计算图像质量得分"""
        if self.adaptive_config['quality_metric'] == 'ssim':
            # 使用SSIM作为质量指标（越高越好）
            ssim_score = self.calculate_ssim(removed, original)
            return ssim_score.item() if hasattr(ssim_score, 'item') else ssim_score
        else:
            # 使用LPIPS作为质量指标（越低越好，需要转换）
            lpips_score = torch.mean(self.lpips_loss(removed, original))
            # 转换为0-1范围，1表示质量最好
            return max(0.0, 1.0 - lpips_score.item())

    def calculate_ssim(self, img1, img2):
        """计算SSIM"""
        # 简化的SSIM计算
        mu1 = torch.mean(img1)
        mu2 = torch.mean(img2)

        sigma1_sq = torch.var(img1)
        sigma2_sq = torch.var(img2)
        sigma12 = torch.mean((img1 - mu1) * (img2 - mu2))

        c1 = 0.01 ** 2
        c2 = 0.03 ** 2

        ssim = ((2 * mu1 * mu2 + c1) * (2 * sigma12 + c2)) / \
               ((mu1 ** 2 + mu2 ** 2 + c1) * (sigma1_sq + sigma2_sq + c2))

        return 1 - ssim  # 转换为损失

    def calculate_watermark_detection_loss(self, removed_images):
        """计算水印检测损失（用于对抗训练）"""
        if self.sepmark_net is None or self.original_message is None:
            return torch.tensor(0.0, device=self.device)  # 确保在GPU上
        
        try:
            with torch.no_grad():
                # 确保所有张量在GPU上
                removed_images = removed_images.to(self.device)
                
                # 使用SepMark解码
                decoded_c = self.sepmark_net.module.decoder_C(removed_images)
                decoded_rf = self.sepmark_net.module.decoder_RF(removed_images)
                
                # 计算错误率（我们希望错误率高，即水印被成功移除）
                batch_size = removed_images.size(0)
                original_message_batch = self.original_message.repeat(batch_size, 1).to(self.device)
                
                # 计算错误率
                error_rate_c = torch.mean(torch.abs(decoded_c - original_message_batch))
                error_rate_rf = torch.mean(torch.abs(decoded_rf - original_message_batch))
                
                # 我们希望错误率高（即水印被移除），所以损失是负的错误率
                adversarial_loss = -(error_rate_c + error_rate_rf)
                
                return adversarial_loss
                
        except Exception as e:
            print(f"水印检测损失计算失败: {e}")
            return torch.tensor(0.0, device=self.device)

    def train_epoch(self, dataloader):
        self.removal_net.train()
        epoch_losses = {
            'total': 0.0,
            'mse': 0.0,
            'lpips': 0.0,
            'ssim': 0.0,
            'adversarial': 0.0
        }
        
        # 动态权重统计
        weight_adjustments = 0
        quality_scores = []

        for batch_idx, (watermarked, original) in enumerate(tqdm(dataloader, desc="训练")):
            watermarked = watermarked.to(self.device)
            original = original.to(self.device)

            # 前向传播
            removed = self.removal_net(watermarked)

            # 计算各项损失
            mse_loss = self.mse_loss(removed, original)
            lpips_loss = torch.mean(self.lpips_loss(removed, original))
            ssim_loss = self.calculate_ssim(removed, original)
            adversarial_loss = self.calculate_watermark_detection_loss(removed)

            # 计算图像质量得分
            quality_score = self.calculate_image_quality_score(removed, original)
            quality_scores.append(quality_score)

            # 使用动态权重计算总损失
            total_loss = (
                    self.dynamic_weights['mse'] * mse_loss +
                    self.dynamic_weights['lpips'] * lpips_loss +
                    self.dynamic_weights['ssim'] * ssim_loss +
                    self.dynamic_weights['adversarial'] * adversarial_loss
            )

            # 动态调整权重（每10个batch检查一次）
            if batch_idx % 10 == 0:
                old_weights = self.dynamic_weights.copy()
                self.adjust_weights_based_on_quality(quality_score, total_loss.item())
                
                # 检查权重是否发生变化
                if old_weights != self.dynamic_weights:
                    weight_adjustments += 1
                    
                    # 重新计算损失（使用新权重）
                    total_loss = (
                            self.dynamic_weights['mse'] * mse_loss +
                            self.dynamic_weights['lpips'] * lpips_loss +
                            self.dynamic_weights['ssim'] * ssim_loss +
                            self.dynamic_weights['adversarial'] * adversarial_loss
                    )

            # 反向传播
            self.optimizer.zero_grad()
            total_loss.backward()
            self.optimizer.step()

            # 记录损失
            epoch_losses['total'] += total_loss.item()
            epoch_losses['mse'] += mse_loss.item()
            epoch_losses['lpips'] += lpips_loss.item()
            epoch_losses['ssim'] += ssim_loss.item()
            epoch_losses['adversarial'] += adversarial_loss.item()

            # 每N个batch打印一次状态
            if batch_idx % self.config.monitoring.print_batch_interval == 0:
                avg_quality = np.mean(quality_scores[-100:]) if len(quality_scores) >= 100 else np.mean(quality_scores)
                print(f"Batch {batch_idx}: 质量={avg_quality:.3f}, 当前权重 MSE={self.dynamic_weights['mse']:.3f}, "
                      f"LPIPS={self.dynamic_weights['lpips']:.3f}, Adversarial={self.dynamic_weights['adversarial']:.3f}")

        # 平均损失
        for key in epoch_losses:
            epoch_losses[key] /= len(dataloader)

        # 添加质量监控信息
        avg_quality = np.mean(quality_scores)
        epoch_losses['avg_quality'] = avg_quality
        epoch_losses['weight_adjustments'] = weight_adjustments
        
        print(f"Epoch总结: 平均质量={avg_quality:.3f}, 权重调整次数={weight_adjustments}")

        return epoch_losses

    def evaluate(self, dataloader):
        """评估模型"""
        self.removal_net.eval()
        eval_losses = {
            'mse': 0.0,
            'lpips': 0.0,
            'watermark_removal_rate': 0.0
        }

        total_samples = 0
        successful_removals = 0

        with torch.no_grad():
            for watermarked, original in dataloader:
                watermarked = watermarked.to(self.device)
                original = original.to(self.device)

                removed = self.removal_net(watermarked)

                # 计算损失
                eval_losses['mse'] += self.mse_loss(removed, original).item()
                eval_losses['lpips'] += torch.mean(self.lpips_loss(removed, original)).item()

                # 评估水印移除效果
                if self.sepmark_net is not None and self.original_message is not None:
                    try:
                        decoded_C = self.sepmark_net.module.decoder_C(removed)

                        orig_binary = (self.original_message > 0).int().cpu().numpy()
                        decoded_binary = (decoded_C > 0).int().cpu().numpy()

                        # 计算错误率（错误率高说明水印移除效果好）
                        for i in range(decoded_binary.shape[0]):
                            error_rate = np.sum(orig_binary != decoded_binary[i]) / len(orig_binary)
                            if error_rate > 0.3:  # 错误率>30%认为移除成功
                                successful_removals += 1
                            total_samples += 1

                    except Exception:
                        pass

        # 平均损失
        for key in ['mse', 'lpips']:
            eval_losses[key] /= len(dataloader)

        if total_samples > 0:
            eval_losses['watermark_removal_rate'] = successful_removals / total_samples

        return eval_losses

    def train(self):
        """主训练循环"""
        print("开始训练去水印网络...")
        
        # 清理GPU内存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            print("GPU内存已清理")

        # 数据变换
        transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

        # 创建数据集
        train_dataset = WatermarkedDataset(
            self.config.data.original_dir,
            self.config.data.watermarked_dir,
            transform=transform,
            max_samples=self.config.training.max_samples
        )

        train_dataloader = DataLoader(
            train_dataset,
            batch_size=self.config.training.batch_size,
            shuffle=self.config.data_loading.shuffle,
            num_workers=self.config.data_loading.num_workers,
            pin_memory=self.config.data_loading.pin_memory
        )

        # 训练循环
        best_loss = float('inf')

        for epoch in range(self.config.training.num_epochs):
            print(f"\nEpoch {epoch + 1}/{self.config.training.num_epochs}")

            # 训练
            train_losses = self.train_epoch(train_dataloader)

            # 调整学习率
            self.scheduler.step()

            # 打印训练结果
            print(f"训练损失:")
            for key, value in train_losses.items():
                print(f"  {key}: {value:.6f}")
            
            # 显示当前权重状态
            print(f"当前权重状态:")
            print(f"  MSE: {self.dynamic_weights['mse']:.4f}")
            print(f"  LPIPS: {self.dynamic_weights['lpips']:.4f}")
            print(f"  SSIM: {self.dynamic_weights['ssim']:.4f}")
            print(f"  Adversarial: {self.dynamic_weights['adversarial']:.4f}")
            
            # 质量警告
            if 'avg_quality' in train_losses:
                if train_losses['avg_quality'] < self.adaptive_config['quality_threshold']:
                    print(f"⚠️  图像质量警告: {train_losses['avg_quality']:.4f} < {self.adaptive_config['quality_threshold']}")
                else:
                    print(f"✓ 图像质量达标: {train_losses['avg_quality']:.4f}")

            # 每10个epoch评估一次
            if (epoch + 1) % 10 == 0:
                eval_losses = self.evaluate(train_dataloader)
                print(f"评估结果:")
                for key, value in eval_losses.items():
                    print(f"  {key}: {value:.6f}")

                # 保存最佳模型
                if eval_losses['mse'] < best_loss:
                    best_loss = eval_losses['mse']
                    self.save_model(f'best_removal_model.pth')
                    print("保存最佳模型")

            # 定期保存检查点
            if (epoch + 1) % self.config.monitoring.save_checkpoint_interval == 0:
                self.save_model(f'removal_model_epoch_{epoch + 1}.pth')
                self.plot_weight_history(epoch + 1)

        print("训练完成！")
        
        # 保存最终的训练历史
        self.save_training_history()

    def save_model(self, filename):
        """保存模型"""
        os.makedirs('checkpoints', exist_ok=True)
        torch.save({
            'model_state_dict': self.removal_net.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'config': self.config,
            'weight_history': self.weight_history,
            'dynamic_weights': self.dynamic_weights
        }, f'checkpoints/{filename}')
    
    def plot_weight_history(self, epoch):
        """绘制权重变化历史"""
        if len(self.weight_history['quality_scores']) < 10:
            return
            
        plt.figure(figsize=(15, 5))
        
        # 质量得分变化
        plt.subplot(1, 3, 1)
        plt.plot(self.weight_history['quality_scores'], label='Quality Score')
        plt.axhline(y=self.adaptive_config['quality_threshold'], color='r', linestyle='--', label='Threshold')
        plt.title('Image Quality Score Over Time')
        plt.xlabel('Batch')
        plt.ylabel('Quality Score')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 权重变化
        plt.subplot(1, 3, 2)
        plt.plot(self.weight_history['quality_weights'], label='Quality Weights (MSE+LPIPS+SSIM)')
        plt.plot(self.weight_history['watermark_weights'], label='Watermark Weights (Adversarial)')
        plt.title('Weight Changes Over Time')
        plt.xlabel('Batch')
        plt.ylabel('Weight Value')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 总损失变化
        plt.subplot(1, 3, 3)
        plt.plot(self.weight_history['total_losses'], label='Total Loss')
        plt.title('Total Loss Over Time')
        plt.xlabel('Batch')
        plt.ylabel('Loss Value')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f'weight_history_epoch_{epoch}.png', dpi=150, bbox_inches='tight')
        plt.close()
        
        print(f"权重历史图已保存: weight_history_epoch_{epoch}.png")
    
    def save_training_history(self):
        """保存训练历史"""
        import json
        
        history_data = {
            'adaptive_config': self.adaptive_config,
            'weight_history': self.weight_history,
            'final_weights': self.dynamic_weights
        }
        
        with open('training_history.json', 'w') as f:
            json.dump(history_data, f, indent=2)
        
        print("训练历史已保存: training_history.json")


def main():
    """主函数"""
    # GPU检查
    print("="*60)
    print("Torch状态检查:", torch.__version__)
    print(f"CUDA可用: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"CUDA设备数量: {torch.cuda.device_count()}")
        print(f"当前设备: {torch.cuda.current_device()}")
        print(f"设备名称: {torch.cuda.get_device_name(0)}")
        print(f"CUDA版本: {torch.version.cuda}")
        
        # 显存检查
        device_id = torch.cuda.current_device()
        memory_allocated = torch.cuda.memory_allocated(device_id) / (1024**3)
        memory_total = torch.cuda.get_device_properties(device_id).total_memory / (1024**3)
        print(f"显存使用: {memory_allocated:.2f} GB / {memory_total:.2f} GB")
    else:
        print("⚠️  警告: CUDA不可用，将使用CPU训练（速度会很慢）")
    print("="*60)
    
    # 创建配置目录
    os.makedirs('config', exist_ok=True)

    # 创建训练器
    trainer = WatermarkRemovalTrainer('config/removal_config.yaml')

    # 开始训练
    trainer.train()


if __name__ == '__main__':
    main()